{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "919e2034",
   "metadata": {},
   "outputs": [],
   "source": [
    "# imports\n",
    "import os\n",
    "import json\n",
    "from dotenv import load_dotenv\n",
    "from openai import OpenAI\n",
    "import gradio as gr\n",
    "\n",
    "import sqlite3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "8330ef87",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "OpenAI API Key exists and begins sk-proj-\n"
     ]
    }
   ],
   "source": [
    "# Initialization\n",
    "load_dotenv(override=True)\n",
    "\n",
    "openai_api_key = os.getenv('OPENAI_API_KEY')\n",
    "if openai_api_key:\n",
    "    print(f\"OpenAI API Key exists and begins {openai_api_key[:8]}\")\n",
    "else:\n",
    "    print(\"OpenAI API Key not set\")\n",
    "    \n",
    "MODEL = \"gpt-4o-mini\"\n",
    "openai = OpenAI()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d5795264",
   "metadata": {},
   "outputs": [],
   "source": [
    "DB = \"prices.db\"\n",
    "\n",
    "with sqlite3.connect(DB) as conn:\n",
    "    cursor = conn.cursor()\n",
    "    cursor.execute('CREATE TABLE IF NOT EXISTS prices (city TEXT PRIMARY KEY, price REAL)')\n",
    "    conn.commit()\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad9cc404",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_ticket_price(city):\n",
    "    print(f\"DATABASE TOOL CALLED: Getting price for {city}\", flush=True)\n",
    "    with sqlite3.connect(DB) as conn:\n",
    "        cursor = conn.cursor()\n",
    "        cursor.execute('SELECT price FROM prices WHERE city = ?', (city.lower(),))\n",
    "        result = cursor.fetchone()\n",
    "        return f\"Ticket price to {city} is ${result[0]}\" if result else \"No price data available for this city\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1bf771c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def set_ticket_price(city, price):\n",
    "    print(f\"DATABASE TOOL CALLED: Setting price for {city} to {price}\", flush=True)\n",
    "    with sqlite3.connect(DB) as conn:\n",
    "        cursor = conn.cursor()\n",
    "        cursor.execute('INSERT INTO prices (city, price) VALUES (?,?) ON CONFLICT(city) DO UPDATE SET price = ?', (city.lower(), price, price))\n",
    "        conn.commit()\n",
    "        return f\"Updated price for {city} to ${price}.\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "852a3a23",
   "metadata": {},
   "outputs": [],
   "source": [
    "ticket_prices = {\"london\":799, \"paris\": 899, \"tokyo\":1400, \"sydney\":2999}\n",
    "for city, price in ticket_prices.items():\n",
    "    set_ticket_price(city, price)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "8a4f0ff0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "DATABASE TOOL CALLED: Getting price for London\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'Ticket price to London is $799.0'"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "get_ticket_price(\"London\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "320c0dcd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# There's a particular dictionary structure that's required to describe our function:\n",
    "get_price_function = {\n",
    "    \"name\": \"get_ticket_price\",\n",
    "    \"description\": \"Get the price of a return ticket to the destination city. Call this whenever you need to know the ticket price, for example when a customer asks 'How much is a ticket to this city'\",\n",
    "    \"parameters\": {\n",
    "        \"type\": \"object\",\n",
    "        \"properties\": {\n",
    "            \"destination_city\": {\n",
    "                \"type\": \"string\",\n",
    "                \"description\": \"The city that the customer wants to travel to\",\n",
    "            },\n",
    "        },\n",
    "        \"required\": [\"destination_city\"],\n",
    "        \"additionalProperties\": False\n",
    "    }\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "545fd6f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "set_price_function = {\n",
    "    \"name\": \"set_ticket_price\",\n",
    "    \"description\": \"Set the price of a return ticket for a given destination. Call this whenever you need to update the ticket price for a city.\",\n",
    "    \"parameters\": {\n",
    "        \"type\": \"object\",\n",
    "        \"properties\": {\n",
    "            \"destination_city\": {\n",
    "                \"type\": \"string\",\n",
    "                \"description\": \"The city whose ticket price needs to be updated\",\n",
    "            },\n",
    "            \"new_price\": {\n",
    "                \"type\": \"number\",\n",
    "                \"description\": \"The return ticket price to assign to the city\",\n",
    "            },\n",
    "        },\n",
    "        \"required\": [\"destination_city\", \"new_price\"],\n",
    "        \"additionalProperties\": False,\n",
    "    },\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "fd369f3d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# And this is included in a list of tools:\n",
    "tools = [{\"type\": \"function\", \"function\": get_price_function},\n",
    "         {\"type\": \"function\", \"function\": set_price_function},]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "9c14c7ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "def handle_tool_calls(message):\n",
    "    responses = []\n",
    "    for tool_call in message.tool_calls:\n",
    "        if tool_call.function.name == \"get_ticket_price\":\n",
    "            arguments = json.loads(tool_call.function.arguments)\n",
    "            city = arguments.get(\"destination_city\")\n",
    "            price_details = get_ticket_price(city)\n",
    "            responses.append({\n",
    "                \"role\": \"tool\",\n",
    "                \"content\": price_details,\n",
    "                \"tool_call_id\": tool_call.id\n",
    "            })\n",
    "        elif tool_call.function.name == \"set_ticket_price\":\n",
    "            arguments = json.loads(tool_call.function.arguments)\n",
    "            city = arguments.get(\"destination_city\")\n",
    "            price = arguments.get(\"new_price\")\n",
    "            confirmation = set_ticket_price(city, price)\n",
    "            responses.append({\n",
    "                \"role\": \"tool\",\n",
    "                \"content\": confirmation,\n",
    "                \"tool_call_id\": tool_call.id\n",
    "            })\n",
    "    return responses"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "858d9575",
   "metadata": {},
   "outputs": [],
   "source": [
    "system_message = \"You are a helpful assistant for an Airline called FlightAI. \"\n",
    "system_message += \"Give short, courteous answers, no more than 1 sentence. \"\n",
    "system_message += \"Always be accurate. If you don't know the answer, say so.\"\n",
    "\n",
    "def chat(user_message, history):\n",
    "    # Normalize history from Gradio format to OpenAI format\n",
    "    history = [{\"role\": h[\"role\"], \"content\": h[\"content\"]} for h in history]\n",
    "\n",
    "    # Start message list with system + history + current user message\n",
    "    messages = (\n",
    "        [{\"role\": \"system\", \"content\": system_message}]\n",
    "        + history\n",
    "        + [{\"role\": \"user\", \"content\": user_message}]\n",
    "    )\n",
    "\n",
    "    # First model call (tools are available)\n",
    "    response = openai.chat.completions.create(\n",
    "        model=MODEL,\n",
    "        messages=messages,\n",
    "        tools=tools,\n",
    "    )\n",
    "\n",
    "    # While the model is asking to call tools\n",
    "    while response.choices[0].finish_reason == \"tool_calls\":\n",
    "        assistant_msg = response.choices[0].message  # assistant message that includes tool_calls\n",
    "\n",
    "        # Let your Python tools run\n",
    "        tool_responses = handle_tool_calls(assistant_msg)\n",
    "\n",
    "        # Add assistant tool-call message and tool outputs to the conversation\n",
    "        messages.append(assistant_msg)\n",
    "        messages.extend(tool_responses)\n",
    "\n",
    "        # Ask the model again, now with tool outputs included\n",
    "        response = openai.chat.completions.create(\n",
    "            model=MODEL,\n",
    "            messages=messages,\n",
    "            tools=tools,\n",
    "        )\n",
    "\n",
    "    # Final assistant message (no more tool calls)\n",
    "    final_msg = response.choices[0].message\n",
    "    return final_msg.content\n",
    "   \n",
    "    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f5e2afd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "* Running on local URL:  http://127.0.0.1:7903\n",
      "* To create a public link, set `share=True` in `launch()`.\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div><iframe src=\"http://127.0.0.1:7903/\" width=\"100%\" height=\"500\" allow=\"autoplay; camera; microphone; clipboard-read; clipboard-write;\" frameborder=\"0\" allowfullscreen></iframe></div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": []
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "DATABASE TOOL CALLED: Getting price for Berlin\n",
      "DATABASE TOOL CALLED: Getting price for London\n"
     ]
    }
   ],
   "source": [
    "gr.ChatInterface(fn=chat, type=\"messages\").launch()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9c46738",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llms_sept25",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
